use ndarray::prelude::*;
use ndarray::{Array1, Array2};
use ndarray_linalg::Solve;
use std::collections::VecDeque;
use std::f64::consts::PI;
use rand::Rng;

// 数据结构：存储时间序列库存数据点
#[derive(Debug, Clone)]
struct DataPoint {
    timestamp: f64, // 时间戳（任意时间单位）
    inventory: f64, // 该时间点的库存量
}

// 趋势+季节性模型
struct TrendSeasonalModel {
    trend_coefficients: Array1<f64>,    // [截距, 斜率]
    seasonal_coefficients: Array1<f64>, // 傅里叶系数
    period: f64,                        // 周期长度
    fourier_terms: usize,               // 傅里叶项数
    last_train_time: f64,               // 最后训练时间
}

// 预测系统配置
struct ForecastConfig {
    max_history: usize,    // 最大历史数据点数
    min_history: usize,    // 最小历史数据点数
    retrain_interval: f64, // 重新训练间隔（时间单位）
    data_decay_rate: f64,  // 数据衰减率
    period: f64,           // 季节性周期
    fourier_terms: usize,  // 傅里叶项数
    min_interval: f64,     // 最小时间间隔阈值
}

// 预测系统状态
struct ForecastingSystem {
    config: ForecastConfig,
    data_buffer: VecDeque<DataPoint>, // 数据缓冲区（按时间排序）
    model: Option<TrendSeasonalModel>, // 当前模型
    last_retrain_time: f64,           // 最后训练时间
}

impl TrendSeasonalModel {
    fn new() -> Self {
        Self {
            trend_coefficients: Array1::zeros(2),
            seasonal_coefficients: Array1::zeros(0),
            period: 12.0,
            fourier_terms: 2,
            last_train_time: 0.0,
        }
    }
    
    // 使用加权最小二乘法拟合模型
    fn fit_weighted(
        times: &Array1<f64>, 
        scrap_rates: &Array1<f64>,  // 现在改为报废率数据
        weights: &Array1<f64>,
        period: f64,
        fourier_terms: usize,
    ) -> Option<Self> {
        let n = times.len();
        let total_terms = 2 + 2 * fourier_terms; // 趋势(2) + 傅里叶(2*terms)
        
        // 调整最小数据要求
        if n < total_terms {
            return None; // 数据不足
        }
        
        let mut features = Array2::zeros((n, total_terms));
        let mut weighted_values = Array1::zeros(n);
        
        // 构建特征矩阵和加权值向量
        for i in 0..n {
            let t = times[i];
            let w = weights[i].sqrt(); // 使用平方根权重进行加权最小二乘
            
            // 趋势特征
            features[[i, 0]] = w * 1.0;      // 常数项
            features[[i, 1]] = w * t;        // 时间项
            
            // 季节性特征（傅里叶级数）
            for k in 1..=fourier_terms {
                let angle = 2.0 * PI * k as f64 * t / period;
                features[[i, 2 + 2*k - 2]] = w * angle.sin(); // sin项
                features[[i, 2 + 2*k - 1]] = w * angle.cos(); // cos项
            }
            
            weighted_values[i] = w * scrap_rates[i];
        }
        
        // 添加正则化项避免奇异矩阵
        let regularization: Array2<f64> = Array2::eye(total_terms) * 1e-6;
        
        // 求解加权最小二乘问题
        let xt = features.t();
        let xtx = xt.dot(&features) + &regularization;
        let xty = xt.dot(&weighted_values);
        
        match xtx.solve(&xty) {
            Ok(all_coefficients) => {
                let trend_coefficients = all_coefficients.slice(s![0..2]).to_owned();
                let seasonal_coefficients = all_coefficients.slice(s![2..]).to_owned();
                
                Some(Self {
                    trend_coefficients,
                    seasonal_coefficients,
                    period,
                    fourier_terms,
                    last_train_time: *times.iter().max_by(|a, b| a.partial_cmp(b).unwrap()).unwrap_or(&0.0),
                })
            },
            Err(e) => {
                eprintln!("模型拟合失败: {:?}", e);
                None
            }
        }
    }
    
    // 预测指定时间点的报废率
    fn predict(&self, time: f64) -> f64 {
        // 趋势成分
        let trend = self.trend_coefficients[0] + self.trend_coefficients[1] * time;
        
        // 季节性成分
        let mut seasonal = 0.0;
        for k in 0..self.fourier_terms {
            let angle = 2.0 * PI * (k + 1) as f64 * time / self.period;
            seasonal += self.seasonal_coefficients[2*k] * angle.sin()
                     + self.seasonal_coefficients[2*k + 1] * angle.cos();
        }
        
        trend + seasonal
    }
    
    // 分解预测值为趋势和季节性成分
    fn decompose(&self, time: f64) -> (f64, f64) {
        let trend = self.trend_coefficients[0] + self.trend_coefficients[1] * time;
        let mut seasonal = 0.0;
        for k in 0..self.fourier_terms {
            let angle = 2.0 * PI * (k + 1) as f64 * time / self.period;
            seasonal += self.seasonal_coefficients[2*k] * angle.sin()
                     + self.seasonal_coefficients[2*k + 1] * angle.cos();
        }
        (trend, seasonal)
    }
}

impl ForecastConfig {
    fn default() -> Self {
        Self {
            max_history: 36,   // 最多保留36个数据点
            min_history: 4,    // 降低最小数据点要求（从6降到4）
            retrain_interval: 1.0, // 重新训练间隔
            data_decay_rate: 0.2,  // 每月衰减20%
            period: 12.0,      // 年度周期
            fourier_terms: 2,  // 使用2个傅里叶项
            min_interval: 0.1, // 最小时间间隔阈值
        }
    }
}

impl ForecastingSystem {
    fn new(config: ForecastConfig) -> Self {
        let max_history = config.max_history;
        Self {
            config,
            data_buffer: VecDeque::with_capacity(max_history),
            model: None,
            last_retrain_time: 0.0,
        }
    }
    
    // 添加新库存数据点（时间间隔不等）
    fn add_data_point(&mut self, timestamp: f64, inventory: f64) {
        // 维护数据缓冲区大小
        if self.data_buffer.len() >= self.config.max_history {
            self.data_buffer.pop_front();
        }
        
        // 插入到正确位置以保持时间顺序
        let index = self.data_buffer.iter()
            .position(|p| p.timestamp > timestamp)
            .unwrap_or(self.data_buffer.len());
        
        self.data_buffer.insert(index, DataPoint { timestamp, inventory });
        
        // 检查是否需要重新训练
        if self.should_retrain(timestamp) {
            self.retrain(timestamp);
        }
    }
    
    // 判断是否需要重新训练
    fn should_retrain(&self, current_time: f64) -> bool {
        if self.data_buffer.len() < self.config.min_history {
            return false; // 数据不足
        }
        
        // 检查是否到达重新训练时间
        current_time - self.last_retrain_time >= self.config.retrain_interval
    }
    
    // 重新训练模型（基于库存变化计算报废率）
    fn retrain(&mut self, current_time: f64) {
        // 需要至少两个点才能计算报废率
        if self.data_buffer.len() < 2 {
            return;
        }
        
        // 准备数据数组
        let n = self.data_buffer.len() - 1;
        let mut times = Vec::with_capacity(n);
        let mut scrap_rates = Vec::with_capacity(n); // 报废率 = 库存变化量 / 时间间隔
        let mut weights = Vec::with_capacity(n);
        
        // 计算每个时间间隔的报废率
        for i in 1..self.data_buffer.len() {
            let prev = &self.data_buffer[i-1];
            let curr = &self.data_buffer[i];
            
            // 计算时间间隔
            let interval = curr.timestamp - prev.timestamp;
            
            // 跳过过小的间隔
            if interval < self.config.min_interval {
                continue;
            }
            
            // 计算库存变化量（报废量）
            let scrap_amount = prev.inventory - curr.inventory;
            
            // 仅处理库存减少的情况（报废量为正）
            if scrap_amount > 0.0 {
                // 计算平均报废率（库存变化量 / 时间间隔）
                let scrap_rate = scrap_amount / interval;
                
                times.push((prev.timestamp + curr.timestamp) / 2.0); // 使用区间中点的时间戳
                scrap_rates.push(scrap_rate);
                
                // 计算权重（基于时间间隔和时效性）
                let time_weight = (-self.config.data_decay_rate * (current_time - curr.timestamp)).exp();
                let interval_weight = interval.sqrt(); // 时间间隔越长，权重越大
                weights.push(time_weight * interval_weight);
            }
        }
        
        // 检查是否有足够的数据点
        if times.len() < self.config.min_history {
            eprintln!("训练数据不足: 需要 {} 个点，实际有 {} 个点", self.config.min_history, times.len());
            return;
        }
        
        // 转换为ndarray类型
        let times_arr = Array1::from(times);
        let scrap_rates_arr = Array1::from(scrap_rates);
        let weights_arr = Array1::from(weights);
        
        // 训练新模型（预测报废率）
        if let Some(new_model) = TrendSeasonalModel::fit_weighted(
            &times_arr,
            &scrap_rates_arr,
            &weights_arr,
            self.config.period,
            self.config.fourier_terms,
        ) {
            println!("模型训练成功！使用 {} 个数据点", times_arr.len());
            self.model = Some(new_model);
            self.last_retrain_time = current_time;
        } else {
            eprintln!("模型训练失败！");
        }
    }
    
    // 预测未来时间点的报废率
    fn predict_scrap_rate(&self, future_time: f64) -> Option<f64> {
        self.model.as_ref().map(|model| model.predict(future_time).max(0.0)) // 确保报废率非负
    }
    
    // 预测多个未来时间点的报废率
    fn predict_scrap_rates(&self, start_time: f64, end_time: f64, step: f64) -> Vec<(f64, f64)> {
        let mut predictions = Vec::new();
        let mut current_time = start_time;
        
        while current_time <= end_time {
            if let Some(prediction) = self.predict_scrap_rate(current_time) {
                predictions.push((current_time, prediction));
            }
            current_time += step;
        }
        
        predictions
    }
    
    // 计算库存使用时间（基于预测的报废率）
    fn estimate_inventory_duration(&self, current_inventory: f64, current_time: f64, future_time: f64) -> f64 {
        if self.data_buffer.is_empty() || current_inventory <= 0.0 {
            return 0.0;
        }
        
        // 获取最后记录的时间点
        let last_point = self.data_buffer.back().unwrap();
        let last_time = last_point.timestamp;
        
        // 计算从最后记录时间到当前时间的库存消耗
        let elapsed_time = current_time - last_time;
        let mut remaining_inventory = current_inventory;
        
        // 如果有模型且时间间隔合理，计算这段时间的消耗
        if elapsed_time > self.config.min_interval {
            if let Some(scrap_rate) = self.predict_scrap_rate(current_time) {
                let consumed = scrap_rate * elapsed_time;
                remaining_inventory = (current_inventory - consumed).max(0.0);
            }
        }
        
        // 如果库存已耗尽
        if remaining_inventory <= 0.0 {
            return elapsed_time; // 已经耗尽
        }
        
        // 计算历史平均报废率作为默认值
        let mut default_scrap_rate = 0.0;
        let mut count = 0;
        
        for i in 1..self.data_buffer.len() {
            let prev = &self.data_buffer[i-1];
            let curr = &self.data_buffer[i];
            let interval = curr.timestamp - prev.timestamp;
            
            if interval > self.config.min_interval {
                let scrap_amount = prev.inventory - curr.inventory;
                if scrap_amount > 0.0 {
                    default_scrap_rate += scrap_amount / interval;
                    count += 1;
                }
            }
        }
        
        if count > 0 {
            default_scrap_rate /= count as f64;
        } else {
            default_scrap_rate = 1000.0; // 默认值
        }
        
        // 预测未来报废率
        let mut cumulative_time = 0.0;
        let step = 0.1; // 精细时间步长
        let mut current = current_time;
        let mut remaining = remaining_inventory;
        
        // 设置最大预测时间（避免无限循环）
        let max_time = current_time + (future_time - current_time) * 2.0;
        
        while current <= max_time && remaining > 0.0 {
            // 使用模型预测的报废率，如果没有则使用历史平均值
            let scrap_rate = self.predict_scrap_rate(current)
                .unwrap_or(default_scrap_rate)
                .max(0.0);
            
            // 计算当前时间步长的消耗
            let consumption = scrap_rate * step;
            
            if consumption >= remaining {
                // 在这个时间步长内库存耗尽
                let fraction = remaining / consumption;
                cumulative_time += fraction * step;
                return elapsed_time + cumulative_time;
            } else {
                // 消耗部分库存
                remaining -= consumption;
                cumulative_time += step;
            }
            current += step;
        }
        
        // 在最大预测时间内库存未耗尽
        elapsed_time + cumulative_time // 返回当前已计算的时间
    }
    
    // 获取模型性能指标
    fn model_performance(&self) -> Option<f64> {
        if self.data_buffer.len() < 6 || self.model.is_none() {
            return None;
        }
        
        let mut total_error = 0.0;
        let mut count = 0;
        
        // 计算历史预测误差（跳过第一个点）
        for i in 1..self.data_buffer.len() {
            let prev = &self.data_buffer[i-1];
            let curr = &self.data_buffer[i];
            let interval = curr.timestamp - prev.timestamp;
            
            if interval < self.config.min_interval {
                continue;
            }
            
            let actual_scrap_rate = (prev.inventory - curr.inventory) / interval;
            if actual_scrap_rate > 0.0 {
                if let Some(predicted_scrap_rate) = self.predict_scrap_rate((prev.timestamp + curr.timestamp) / 2.0) {
                    let error = (predicted_scrap_rate - actual_scrap_rate).abs() / actual_scrap_rate.max(1.0);
                    total_error += error;
                    count += 1;
                }
            }
        }
        
        if count > 0 {
            Some(total_error / count as f64)
        } else {
            None
        }
    }
}

// 示例使用
fn main() {
    // 创建预测系统
    let mut config = ForecastConfig::default();
    config.min_history = 3; // 降低最小历史数据要求
    config.fourier_terms = 1; // 减少傅里叶项数
    let mut system = ForecastingSystem::new(config);
    
    // 模拟不等间隔的库存数据（添加更多数据点）
    let mut current_time = 0.0;
    let mut inventory = 15000.0; // 初始库存
    let mut rng = rand::thread_rng();
    
    // 添加初始数据点
    system.add_data_point(current_time, inventory);
    println!("添加初始库存: {:.1} at {:.1}", inventory, current_time);
    
    // 模拟12个月的库存变化
    for _ in 1..=12 {
        // 随机间隔（0.5-1.5个月）
        let interval = rng.gen_range(0.5..1.5);
        current_time += interval;
        
        // 随机报废量（800-1200）
        let scrap_amount = rng.gen_range(800.0..1200.0);
        inventory -= scrap_amount;
        
        system.add_data_point(current_time, inventory);
        println!("添加库存点: {:.1} at {:.1} (报废: {:.1})", inventory, current_time, scrap_amount);
        
        // 显式训练模型
        system.retrain(current_time);
    }
    
    // 预测未来报废率
    println!("\n未来报废率预测:");

    for month in 1..=6 {
        let future_time = current_time + month as f64;
        if let Some(rate) = system.predict_scrap_rate(future_time) {
            println!("  {:.1}个月后: {:.1}件/月", month, rate);
        } else {
            println!("  {:.1}个月后: 无法预测", month);
        }
    }
    
    // 计算库存使用时间
    let current_inventory = 5000.0;
    let duration = system.estimate_inventory_duration(current_inventory, current_time, current_time + 120.0);
    
    if duration > 120.0 {
        println!("\n当前库存 {} 件预计还能使用超过 {:.1} 个月", current_inventory, 120.0);
    } else {
        println!("\n当前库存 {} 件预计还能使用 {:.1} 个月", current_inventory, duration);
    }
    
    // 检查模型性能
    if let Some(mape) = system.model_performance() {
        println!("\n模型平均绝对百分比误差: {:.2}%", mape * 100.0);
    } else {
        println!("\n无法计算模型性能");
    }
}